
package org.apache.solr.handler.component;

import java.net.MalformedURLException;
import java.util.Random;
import java.util.concurrent.*;
import org.apache.http.client.HttpClient;
import org.apache.solr.client.solrj.impl.HttpClientUtil;
import org.apache.solr.client.solrj.impl.LBHttpSolrServer;
import org.apache.solr.common.SolrException;
import org.apache.solr.common.params.ModifiableSolrParams;
import org.apache.solr.common.util.ExecutorUtil;
import org.apache.solr.common.util.NamedList;
import org.apache.solr.core.PluginInfo;
import org.apache.solr.util.DefaultSolrThreadFactory;
import org.apache.solr.util.plugin.PluginInfoInitialized;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

public class HttpShardHandlerFactory extends ShardHandlerFactory implements PluginInfoInitialized {

    protected static Logger log = LoggerFactory.getLogger(HttpShardHandlerFactory.class);
    // We want an executor that doesn't take up any resources if
    // it's not used, so it could be created statically for
    // the distributed search component if desired.
    //
    // Consider CallerRuns policy and a lower max threads to throttle
    // requests at some point (or should we simply return failure?)
    ThreadPoolExecutor commExecutor = new ThreadPoolExecutor(
            0,
            Integer.MAX_VALUE,
            5, TimeUnit.SECONDS, // terminate idle threads after 5 sec
            new SynchronousQueue<Runnable>(), // directly hand off tasks
            new DefaultSolrThreadFactory("httpShardExecutor"));

    private HttpClient defaultClient;
    LBHttpSolrServer loadbalancer;
    //default values:
    int soTimeout = 0;
    int connectionTimeout = 0;
    int maxConnectionsPerHost = 20;
    int corePoolSize = 0;
    int maximumPoolSize = Integer.MAX_VALUE;
    int keepAliveTime = 5;
    int queueSize = -1;
    boolean accessPolicy = false;
    public String scheme = "http://"; //current default values
    final Random r = new Random();
    // URL scheme to be used in distributed search.
    static final String INIT_URL_SCHEME = "urlScheme";
    // The core size of the threadpool servicing requests
    static final String INIT_CORE_POOL_SIZE = "corePoolSize";
    // The maximum size of the threadpool servicing requests
    static final String INIT_MAX_POOL_SIZE = "maximumPoolSize";
    // The amount of time idle threads persist for in the queue, before being killed
    static final String MAX_THREAD_IDLE_TIME = "maxThreadIdleTime";
    // If the threadpool uses a backing queue, what is its maximum size (-1) to use direct handoff
    static final String INIT_SIZE_OF_QUEUE = "sizeOfQueue";
    // Configure if the threadpool favours fairness over throughput
    static final String INIT_FAIRNESS_POLICY = "fairnessPolicy";

    /**
     * Get {@link ShardHandler} that uses the default http client.
     */
    @Override
    public ShardHandler getShardHandler() {
        return getShardHandler(defaultClient);
    }

    /**
     * Get {@link ShardHandler} that uses custom http client.
     */
    public ShardHandler getShardHandler(final HttpClient httpClient) {
        return new HttpShardHandler(this, httpClient);
    }

    @Override
    public void init(PluginInfo info) {

        NamedList args = info.initArgs;
        this.soTimeout = getParameter(args, HttpClientUtil.PROP_SO_TIMEOUT, soTimeout);
        this.scheme = getParameter(args, INIT_URL_SCHEME, "http://");
        this.scheme = (this.scheme.endsWith("://")) ? this.scheme : this.scheme + "://";
        this.connectionTimeout = getParameter(args, HttpClientUtil.PROP_CONNECTION_TIMEOUT, connectionTimeout);
        this.maxConnectionsPerHost = getParameter(args, HttpClientUtil.PROP_MAX_CONNECTIONS_PER_HOST, maxConnectionsPerHost);
        this.corePoolSize = getParameter(args, INIT_CORE_POOL_SIZE, corePoolSize);
        this.maximumPoolSize = getParameter(args, INIT_MAX_POOL_SIZE, maximumPoolSize);
        this.keepAliveTime = getParameter(args, MAX_THREAD_IDLE_TIME, keepAliveTime);
        this.queueSize = getParameter(args, INIT_SIZE_OF_QUEUE, queueSize);
        this.accessPolicy = getParameter(args, INIT_FAIRNESS_POLICY, accessPolicy);

        BlockingQueue<Runnable> blockingQueue = (this.queueSize == -1)
                                                ? new SynchronousQueue<Runnable>(this.accessPolicy)
                                                : new ArrayBlockingQueue<Runnable>(this.queueSize, this.accessPolicy);

        this.commExecutor = new ThreadPoolExecutor(
                this.corePoolSize,
                this.maximumPoolSize,
                this.keepAliveTime, TimeUnit.SECONDS,
                blockingQueue,
                new DefaultSolrThreadFactory("httpShardExecutor"));

        ModifiableSolrParams clientParams = new ModifiableSolrParams();
        clientParams.set(HttpClientUtil.PROP_MAX_CONNECTIONS_PER_HOST, maxConnectionsPerHost);
        clientParams.set(HttpClientUtil.PROP_MAX_CONNECTIONS, 10000);
        clientParams.set(HttpClientUtil.PROP_SO_TIMEOUT, soTimeout);
        clientParams.set(HttpClientUtil.PROP_CONNECTION_TIMEOUT, connectionTimeout);
        clientParams.set(HttpClientUtil.PROP_USE_RETRY, false);
        this.defaultClient = HttpClientUtil.createClient(clientParams);

        try {
            loadbalancer = new LBHttpSolrServer(defaultClient);
        }
        catch(MalformedURLException e) {
            // should be impossible since we're not passing any URLs here
            throw new SolrException(SolrException.ErrorCode.SERVER_ERROR, e);
        }

    }

    private <T> T getParameter(NamedList initArgs, String configKey, T defaultValue) {

        T toReturn = defaultValue;
        if(initArgs != null) {
            T temp = (T)initArgs.get(configKey);
            toReturn = (temp != null) ? temp : defaultValue;
        }
        log.info("Setting {} to: {}", configKey, toReturn);
        return toReturn;
    }

    @Override
    public void close() {

        try {
            defaultClient.getConnectionManager().shutdown();
        }
        catch(Throwable e) {
            SolrException.log(log, e);
        }
        try {
            loadbalancer.shutdown();
        }
        catch(Throwable e) {
            SolrException.log(log, e);
        }
        try {
            ExecutorUtil.shutdownNowAndAwaitTermination(commExecutor);
        }
        catch(Throwable e) {
            SolrException.log(log, e);
        }
    }
}
